#include "main.h"

TreeNode* _insert(int key, int data)
{
	if (g_root == NULL)
	{
		g_root = new TreeNode;
		g_root->color = TreeNode::Black;
	}

	TreeNode* p = g_root;
	TreeNode* q = NULL;

	while (p->leftChild != NULL || p->rightChild != NULL)
	{
		if (key < p->key)
		{
			q = p;
			p = p->leftChild;
		}
		else
		{
			q = p;
			p = p->rightChild;
		}
	}

	p->key = key;
	p->data = data;
	p->color = TreeNode::Red;
	p->leftChild = new TreeNode;
	p->leftChild->color = TreeNode::Black;
	p->rightChild = new TreeNode;
	p->rightChild->color = TreeNode::Black;
	p->parent = q;

	return p;
}

inline bool isCollide(TreeNode* node)
{
	if (node->parent == NULL)
	{
		return false;
	}

	if (node->color == TreeNode::Red && node->parent->color == TreeNode::Red)
	{
		return true;
	}

	return false;
}

TreeNode* remapColor(TreeNode* node)
{
	if (node->parent == NULL)
	{
		node->color = TreeNode::Black;
		return node;
	}
	
	if (node->parent->parent == NULL)
	{
		return node;
	}

	if (!isCollide(node))
	{
		return node;
	}

	TreeNode* grandpa = node->parent->parent;

	if (grandpa->leftChild->color == TreeNode::Red && grandpa->rightChild->color == TreeNode::Red)
	{
		grandpa->leftChild->color = TreeNode::Black;
		grandpa->rightChild->color = TreeNode::Black;
		grandpa->color = TreeNode::Red;
		return remapColor(grandpa);
	}
	else
	{
		return node;
	}
}

void rotate(TreeNode* node, bool left)
{
	TreeNode* node2 = NULL;
	TreeNode* tmp = NULL;

	if (left)
	{
		node2 = node->rightChild;
	}
	else
	{
		node2 = node->leftChild;
	}

	node2->parent = node->parent;

	if (node->parent)
	{
		if (node->parent->leftChild == node)
		{
			node->parent->leftChild = node2;
		}
		else
		{
			node->parent->rightChild = node2;
		}
	}

	node->parent = node2;

	if (left)
	{
		tmp = node2->leftChild;
		node2->leftChild = node;
		node->rightChild = tmp;
	}
	else
	{
		tmp = node2->rightChild;
		node2->rightChild = node;
		node->leftChild = tmp;
	}

	if (tmp)
	{
		tmp->parent = node;
	}
}

void rb_insert(int key, int data)
{
	TreeNode* node = _insert(key, data);
	node = remapColor(node);

	if (!isCollide(node))
	{
		return;
	}

	TreeNode* p = node->parent;

	if (p->parent == NULL)
	{
		p->color = TreeNode::Black;
		return;
	}


	if (node == p->leftChild && p == p->parent->rightChild)
	{
		rotate(p, false);
		node = p;
		p = p->parent;
	}
	else if (node == p->rightChild && p == p->parent->leftChild)
	{
		rotate(p, true);
		node = p;
		p = p->parent;
	}
	
	if (node == p->leftChild && p == p->parent->leftChild)
	{
		rotate(p->parent, false);
		p->color = TreeNode::Black;
		p->rightChild->color = TreeNode::Red;
	}
	else
	{
		rotate(p->parent, true);
		p->color = TreeNode::Black;
		p->leftChild->color = TreeNode::Red;
	}
	
	if (p->parent == NULL)
	{
		g_root = p;
	}
}
